package com.dong.redBlackTree;

import com.dong.redBlackTree.entity.RBNode;

import static com.dong.redBlackTree.entity.NodeColor.BLACK;
import static com.dong.redBlackTree.entity.NodeColor.RED;

/**
 * 红黑树
 * 红黑树的主要规则:
 * 1.每个节点不是红色就是黑色的；
 * 2.根节点总是黑色的；
 * 3.如果节点是红色的，则它的子节点必须是黑色的（反之不一定）；
 * 4.从根节点到叶节点或空子节点的每条路径，必须包含相同数目的黑色节点（即相同的黑色高度）;
 */
public class TreeMap<T extends Comparable<T>> {
    private RBNode<T> root;// 根节点

    /**
     * 插入元素
     */
    public void insert(T key) {
        /*
         * 在红-黑树中插入的节点都是红色的，这不是偶然的，
         * 因为插入一个红色节点比插入一个黑色节点违背红-黑规则的可能性更小。
         * 原因是：插入黑色节点总会改变黑色高度（违背规则4），但是插入红色节点只有一半的机会会违背规则3。
         * 另外违背规则3比违背规则4要更容易修正。
         */
        RBNode<T> node = new RBNode<T>(key, RED, null, null, null);
        if (node != null) {
            this.insert(node);
        }
    }

    /**
     * 向红黑树中插入节点
     */
    private void insert(RBNode<T> node){
        RBNode<T> current = null;// 表示最后node的父节点
        RBNode<T> x = this.root;// 用来向下搜索用的
        // 1.找到插入的位置
        while(x != null){
            current = x;
            int cmp = node.getKey().compareTo(x.getKey());
            if(cmp < 0){
                x = x.getLeft();
            }else {
                x= x.getRight();
            }
        }
        // 找到了位置，将当前current作为node的父节点
        node.setParent(current);
        // 2.接下来判断node是插在左子节点还是右子节点
        if (current != null) {
            int cmp = node.getKey().compareTo(current.getKey());
            if (cmp < 0) {
                current.setLeft(node);
            } else {
                current.setRight(node);
            }
        } else {
            this.root = node;
        }
        //3.将它重新修整为一颗红黑树
        this.insertFixUp(node);
    }

    /**
     * 重新修整为红黑树
     */
    private void insertFixUp(RBNode<T> node){
        // 定义父节点和祖父节点
        RBNode<T> parent, gparent;
        // 需要修整的条件：父节点存在，且父节点的颜色是红色
        while ((parent = node.getParent()) != null && this.isRed(parent)) {
            // 获得祖父节点（父节点是红色，肯定有祖父节点）
            gparent = parent.getParent();
            // 若父节点是祖父节点的左子节点，下面else与其相反
            if (parent == gparent.getLeft()) {
                // 获得叔叔节点
                RBNode<T> uncle = gparent.getRight();
                // case1: 叔叔节点也是红色
                if (uncle != null && this.isRed(uncle)) {
                    this.setBlack(parent);
                    this.setBlack(uncle);
                    this.setRed(gparent);
                    node = gparent;
                    continue;// 继续while，重新判断
                }
                // case2: 叔叔节点是黑色，且当前节点是右子节点
                if (node == parent.getRight()) {
                    this.leftRotate(parent);// 从父节点处左旋
                    RBNode<T> tmp = parent;// 然后将父节点和自己调换一下，为下面右旋做准备
                    parent = node;
                    node = tmp;
                }
                // case3: 叔叔节点是黑色，且当前节点是左子节点
                this.setBlack(parent);
                this.setRed(gparent);
                this.rightRotate(gparent);
            } else {//若父节点是祖父节点的右子节点,与上面的完全相反，本质一样的
                // 获得叔叔节点
                RBNode<T> uncle = gparent.getLeft();
                // case1: 叔叔节点也是红色
                if (uncle != null && this.isRed(uncle)) {
                    this.setBlack(parent);
                    this.setBlack(uncle);
                    this.setRed(gparent);
                    node = gparent;
                    continue;// 继续while，重新判断
                }
                // case2: 叔叔节点是黑色，且当前节点是左子节点
                if (parent == gparent.getLeft()) {
                    this.rightRotate(parent);// 从父节点处右旋
                    RBNode<T> tmp = parent;// 然后将父节点和自己调换一下，为下面左旋做准备
                    parent = node;
                    node = tmp;
                }
                // case3: 叔叔节点是黑色，且当前节点是右子节点
                this.setBlack(parent);
                this.setRed(gparent);
                this.leftRotate(gparent);
            }
        }
        // 将根节点设为黑色
        this.setBlack(this.root);
    }

    /**
     * 对节点x进行左旋操作
     * 左旋示意图：
     *        p                       p
     *       /                       /
     *      x                       y
     *     / \                     / \
     *    lx  y      ----->       x  ry
     *       / \                 / \
     *      ly ry               lx ly
     * 左旋做了三件事：
     * 1.将y的左子节点赋给x的右子节点，并将x赋给y左子节点的父节点（y的左子节点非空时）
     * 2.将x的父节点p（非空时）赋给y的父节点，同事更新p的子节点为y（左或右）
     * 3.将y的左子节点设为x，并将x的父节点设为y
     */
    private void leftRotate(RBNode<T> x) {
        RBNode<T> y = x.getRight();
        // step 1
        x.setRight(y.getLeft());
        if (y.getLeft() != null) {
            y.getLeft().setParent(x);
        }
        // step 2
        y.setParent(x.getParent());
        if (x.getParent() == null) {
            this.root = y;
        } else {
            if (x == x.getParent().getLeft()) {
                x.getParent().setLeft(y);
            } else {
                x.getParent().setRight(y);
            }
        }
        // step 3
        y.setLeft(x);
        x.setParent(y);
    }

    /**
     * 对节点y进行右旋操作
     * 右旋示意图：
     *          p                   p
     *         /                   /
     *        y                   x
     *       / \                 / \
     *      x  ry   ----->      lx  y
     *     / \                     / \
     *    lx  rx                   rx ry
     * 右旋做了三件事：
     * 1.将x的右子节点赋给y的左子节点，并将x的右子节点的父节点设为y（x的右子节点非空时）
     * 2.将y的父节点p（非空时）赋给x的父节点，同事更新p的子节点为x（左或右）
     * 3.将x的右节点设为y，并将y的父节点设为x
     */
    private void rightRotate(RBNode<T> y) {
        RBNode<T> x = y.getLeft();
        // step 1
        y.setLeft(x.getRight());;
        if (x.getRight() != null) {
            x.getRight().setParent(y);
        }
        // step 2
        x.setParent(y.getParent());
        if (y.getParent() == null) {
            this.root = x;
        } else {
            if (y == y.getParent().getLeft()) {
                y.getParent().setLeft(x);
            } else {
                y.getParent().setRight(x);
            }
        }
        // step 3
        x.setRight(y);
        y.setParent(x);
    }

    private boolean isRed(RBNode<T> node){
        return node.isColor() == RED;
    }

    private void setRed(RBNode<T> node){
        node.setColor(RED);
    }

    private void setBlack(RBNode<T> node){
        node.setColor(BLACK);
    }
}
